import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.memory import ConversationBufferMemory
from langchain.chains import LLMChain

# --- 环境准备 ---
print("正在从 env.example 加载环境变量...")
load_dotenv(dotenv_path="env.example")
api_key = os.getenv("OPENAI_API_KEY")
model_name = os.getenv("OPENAI_MODEL", "qwen-plus-latest")

if not api_key:
    print("错误：请先在 env.example 文件中设置你的 OPENAI_API_KEY。")
else:
    print("环境变量加载成功！")
    # --- 核心代码开始 ---

    # 1. 初始化模型
    llm = ChatOpenAI(model=model_name)
    print(f"模型已初始化 (使用 {model_name})。")

    # 2. 定义 Memory
    # `memory_key="chat_history"`: 这个名字很重要，它必须和下面提示模板中的占位符名称对应。
    # `return_messages=True`: 我们希望它返回的是消息对象列表，而不是一个大字符串。
    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
    print("记忆模块已初始化。")

    # 3. 定义带"记忆"的提示模板
    # 这次模板有一个新的占位符：MessagesPlaceholder
    # 它专门用于插入一个或多个消息对象。
    prompt = ChatPromptTemplate(
        input_variables=["input"],
        messages=[
            MessagesPlaceholder(variable_name="chat_history"), # 记忆模块存储的历史记录会在这里被插入
            ("user", "{input}") # 用户的当前输入
        ]
    )
    print("带记忆的提示模板已创建。")

    # 4. 创建一个传统的 LLMChain
    # 注意：当处理需要 Memory 的链时，使用 `LLMChain` 会更直接方便。
    # 它专门设计用于处理这种有状态的对话场景。
    chain = LLMChain(
        llm=llm,
        prompt=prompt,
        memory=memory,
        verbose=True  # 设置为 True，可以看到链的详细执行过程，有助于理解！
    )
    print("LLMChain 已创建。")

    # 5. 开始对话！
    print("\n--- 开始与 AI 对话 (输入 '退出' 来结束) ---")
    
    # 第一次对话
    print("\n[你]: 你好，我叫小明。")
    response_1 = chain.invoke({"input": "你好，我叫小明。"})
    print(f"[AI]: {response_1['text']}")

    # 第二次对话
    print("\n[你]: 我叫什么名字？")
    response_2 = chain.invoke({"input": "我叫什么名字？"})
    print(f"[AI]: {response_2['text']}")

    # 第三次对话 - 验证它是否真的记住了
    print("\n[你]: 我今天心情不错，给我讲个关于编程的笑话吧。")
    response_3 = chain.invoke({"input": "我今天心情不错，给我讲个关于编程的笑话吧。"})
    print(f"[AI]: {response_3['text']}")
    
    print("\n--- 对话结束 ---")

    # 你可以手动查看 memory 中存储的内容
    print("\n--- 查看 Memory 中的历史记录 ---")
    print(memory.load_memory_variables({}))
    print("--------------------") 